module time_stepping_mod

  use constants_mod
  use namelist_mod
  use boundary_mod
  use butcher_tableau_mod
  use weno_mod
  use upwind_mod
  use finite_diff_mod
  use MCV_mod
  use TVD_mod

  implicit none
  
  private
  public time_stepping_init
  public time_stepping_final
  public time_integrator

  real(r8), allocatable, dimension(:,:  ) :: matrixA
  real(r8), allocatable, dimension(:    ) :: vectorB
  real(r8), allocatable, dimension(:,:,:) :: Jac
  real(r8), allocatable, dimension(:,:  ) :: correction
  real(r8), allocatable, dimension(:,:  ) :: stage_value
  real(r8), allocatable, dimension(:,:  ) :: tendency

  interface
    subroutine time_integrator_interface(q, dt)
      import r8, nx, nhalo
      real(r8), intent(inout) :: q(1-nhalo:nx+nhalo)
      real(r8), intent(in) :: dt
    end subroutine time_integrator_interface
    subroutine compute_flux_jacobian_interface(q, dt, jac)
      import r8, nx, nhalo
      real(r8), intent(in) :: q(1-nhalo:nx+nhalo)
      real(r8), intent(in) :: dt
      real(r8), intent(out) :: jac(1:nx, 1:nx)
    end subroutine compute_flux_jacobian_interface
    subroutine compute_interface_flux_interface(q, dt, flux)
      import r8, nx, nhalo
      real(r8), intent(in) :: q(1-nhalo:nx+nhalo)
      real(r8), intent(in) :: dt
      real(r8), intent(out) :: flux(0:nx)
    end subroutine compute_interface_flux_interface
    function ode_rhs_interface(q, dt) result(rhs)
      import r8, nx, nhalo
      real(r8), intent(in) :: q(1-nhalo:nx+nhalo)
      real(r8), intent(in) :: dt
      real(r8) :: rhs(1:nx)
    end function ode_rhs_interface
  end interface
  procedure(time_integrator_interface       ), pointer :: time_integrator
  procedure(compute_flux_jacobian_interface ), pointer :: compute_flux_jacobian
  procedure(compute_interface_flux_interface), pointer :: compute_interface_flux
  procedure(ode_rhs_interface               ), pointer :: ode_rhs
  
  contains

  subroutine time_stepping_init()

    if (.not. allocated(matrixA    )) allocate(matrixA    (nx,nx))
    if (.not. allocated(vectorB    )) allocate(vectorB    (nx   ))
    if (.not. allocated(Jac        )) allocate(Jac        (nx,nx,nstage))
    if (.not. allocated(correction )) allocate(correction (nx,nstage))
    if (.not. allocated(stage_value)) allocate(stage_value(1-nhalo:nx+nhalo,nstage))
    if (.not. allocated(tendency   )) allocate(tendency   (1:nx,nstage))

    select case (trim(time_scheme))
    case ('ERK_FE', 'ERK3_WRF', 'ERK4')
      time_integrator => ERK
    case ('Backward_Euler', 'implicit_midpoint', 'Crank_Nicolson', 'ssp2', 'ssp3')
      time_integrator => IRK
    case default
      print*, 'Unknown time_scheme:', trim(time_scheme)
      stop
    end select

    if (trim(space_scheme) == 'MCV') then
      ode_rhs => ode_rhs_mcv
    else
      ode_rhs => ode_rhs_normal
    end if

    select case (trim(space_scheme))
    case ('finite_diff_2nd', 'finite_diff_4th')
      compute_flux_jacobian => compute_flux_jacobian_finite_diff
      compute_interface_flux => compute_interface_flux_finite_diff
    case ('upwind1', 'upwind3')
      compute_flux_jacobian => compute_flux_jacobian_upwind
      compute_interface_flux => compute_interface_flux_upwind
    case ('weno5')
      compute_flux_jacobian => compute_flux_jacobian_weno5
      compute_interface_flux => compute_interface_flux_weno5
    case ('MCV')
      compute_flux_jacobian => compute_flux_jacobian_MCV
    case ('TVD')
      compute_flux_jacobian => compute_flux_jacobian_TVD
      compute_interface_flux => compute_interface_flux_TVD
    case default
      print*, 'Unknown space_scheme: ', trim(space_scheme)
      stop
    end select

  end subroutine time_stepping_init

  subroutine time_stepping_final()

    if (allocated(matrixA    )) deallocate(matrixA   )
    if (allocated(vectorB    )) deallocate(vectorB   )
    if (allocated(Jac        )) deallocate(Jac       )
    if (allocated(correction )) deallocate(correction)
    if (allocated(stage_value)) deallocate(stage_value)
    if (allocated(tendency   )) deallocate(tendency  )

  end subroutine time_stepping_final

  subroutine ERK(var, dt)

    real(r8), intent(inout) :: var(1-nhalo:nx+nhalo)
    real(r8), intent(in) :: dt
    integer istage, j

    tendency(1:nx,1) = ode_rhs(var, dt)
    do istage = 2, nstage
      stage_value(1:nx,istage) = var(1:nx)
      do j = 1, istage - 1
        stage_value(1:nx,istage) = stage_value(1:nx,istage) + ERK_a(istage, j) * dt * tendency(1:nx,j)
      end do
      call set_periodic_boundary(stage_value(1-nhalo:nx+nhalo,istage))
      tendency(1:nx,istage) = ode_rhs(stage_value(1-nhalo:nx+nhalo,istage), dt)
    end do

    do istage = 1, nstage
      var(1:nx) = var(1:nx) + dt * (ERK_b(istage) * tendency(1:nx,istage))
    end do

    call set_periodic_boundary(var)

  end subroutine ERK

  subroutine IRK(var, dt)
    ! =================================================== 
    ! ∂y/∂t = f(y)
    ! Yi = yn +  dt \sum_{j=1}^s Aij * f(Yj) (j=1,2,...,s) 
    ! yn+1 = yn + h  \sum_{i=1}^s Bi * f(Yi) (i=1,2,...,s)
    ! stage solution: G(Yi) = Yi - yn - dt \sum_{j=1}^s Aij * f(Yj)
    ! (I - dt A11 J1)𝝳Y1 = -G(Y1) = yn - Y1 + dt A11 * f(Y1)
    ! (I - dt A22 J2)𝝳Y2 = -G(Y2) + dt A11 J1 𝝳Y1 = yn - Y2 + dt A21f(Y1) + dt A22 f(Y2)
    ! J = d/dy f(y)
    ! =================================================== 
    real(r8), intent(inout) :: var(1-nhalo:nx+nhalo)
    real(r8), intent(in) :: dt
    real(r8) :: err
    integer :: iter, istage, i
    
    jac = 0.0_r8
    do istage = 1, nstage   
      stage_value(1-nhalo:nx+nhalo,istage) = var
      err = huge(1.0_r8)
      iter = 0
      do while (err > tol .and. iter < max_iter)
        call compute_flux_jacobian(stage_value(1-nhalo:nx+nhalo,istage), dt, Jac(1:nx, 1:nx, istage))
        ! call print_sparsity_pattern(jac(1:nx,1:nx, 1))
        ! call output_jacobian(Jac(1:nx,1:nx, istage))
        matrixA = -dt * IRK_a(istage, istage) * Jac(1:nx, 1:nx, istage)
        forall (i=1:nx) matrixA(i,i) = 1.0_r8 + matrixA(i,i)
        ! print*, "IRK stage =", istage, "matrixA =", matrixA(1:nx,1:nx)
        call compute_RK_rhs(istage, var, stage_value(1-nhalo:nx+nhalo, 1:istage), correction(1:nx,1:istage), Jac(1:nx,1:nx,1:istage), dt, vectorB)
        correction(1:nx,istage) = solve_linear_system(matrixA, vectorB, nx)

        stage_value(1:nx,istage) = stage_value(1:nx,istage) + correction(1:nx,istage)
        call set_periodic_boundary(stage_value(1-nhalo:nx+nhalo,istage))
        err = norm2(correction(1:nx,istage)) !/ (norm2(stage_value(1:nx,istage)) + SMALL)
        print*, "Iteration =", iter, "Error =", err
        iter = iter + 1
      end do
    end do

    do istage = 1, nstage
      var(1:nx) = var(1:nx) + dt * IRK_b(istage) * ode_rhs(stage_value(1-nhalo:nx+nhalo,istage), dt)
    end do
 
    call set_periodic_boundary(var)

  end subroutine IRK

  subroutine compute_RK_rhs(mstage, old_solution, stage_solution, deltaY, Jac_matrix, dt, rhs)
    ! (I - dt A11 J1)𝝳Y1 = -G(Y1) = yn - Y1 + dt A11 * f(Y1)
    ! (I - dt A22 J2)𝝳Y2 = -G(Y2) + dt A21 J1 𝝳Y1 = yn - Y2 + dt A21f(Y1) + dt A22 f(Y2) + dt A21 J1 𝝳Y1
    integer, intent(in) :: mstage
    real(r8), intent(in) :: old_solution(1-nhalo:nx+nhalo), &
                          stage_solution(1-nhalo:nx+nhalo, 1:mstage),&
                                  deltaY(1:nx, 1:mstage), &
                          Jac_matrix(1:nx, 1:nx, 1:mstage), dt
    real(r8), intent(out) :: rhs(1:nx)
    integer j

    rhs = old_solution(1:nx) - stage_solution(1:nx, mstage)
    
    do j = 1, mstage
      rhs = rhs + dt * IRK_a(mstage, j) * ode_rhs(stage_solution(1-nhalo:nx+nhalo,j), dt)
      if (j < mstage) then
        rhs = rhs + dt * IRK_a(mstage, j) * matmul(Jac_matrix(1:nx,1:nx,j), deltaY(1:nx,j))
      end if
    end do      

  end subroutine compute_RK_rhs

  function ode_rhs_normal(var, dt) result(rhs)

    real(r8), intent(in) :: var(1-nhalo:nx+nhalo)
    real(r8), intent(in) :: dt
    real(r8) :: rhs(1:nx), flux(0:nx)

    call compute_interface_flux(var, dt, flux)
    rhs(1:nx) = -(flux(1:nx) - flux(0:nx-1)) / dx

  end function ode_rhs_normal

  function solve_linear_system(coefficient_matrix, rhs_vector, dimension_size) result(solution_vector)
    ! input: 
    !   coefficient_matix - dimension_size × dimensizon_size coefficient matrix
    !   rhs_vector - dimension_size dimension vector of right-hand side
    !   dimension_size - dimension size of the matrix
    ! output:
    !   solution_vector - dimension_size dimension vector of solution
    ! 使用LAPACK的DGESV求解器, general solver, works for any matrix
    integer, intent(in) :: dimension_size
    real(r8), intent(in) :: coefficient_matrix(dimension_size,dimension_size), rhs_vector(dimension_size)
    real(r8) :: solution_vector(dimension_size)
    integer :: info, ipiv(dimension_size)

    solution_vector = rhs_vector
    call dgesv(dimension_size, 1, coefficient_matrix, dimension_size, ipiv, solution_vector, dimension_size, info)
    if (info /= 0) then
        print*, "Error: dgesv returned info =", info
        stop
    end if

  end function solve_linear_system

  subroutine print_sparsity_pattern(Jac, threshold)
    real(r8), intent(in) :: Jac(:, :)         ! Jacobian matrix
    real(r8), optional, intent(in) :: threshold       ! Threshold to ignore small values
    integer :: i, j, N

    N = size(Jac, 1)  ! Assume square matrix

    ! Print column indices (header)
    write(*, '(A8)', advance='no') 'Row\Col'
    do j = 1, N
        write(*, '(i5)', advance='no') j
    end do
    print *

    ! Print each row with '*' for non-zero entries
    do i = 1, N
        write(*, '(I5, A1)', advance='no') i, '|'
        do j = 1, N
            ! if (abs(Jac(i, j)) > threshold) then
            !     write(*, '(3x,A3)', advance='no') ' * '
            ! else
            !     write(*, '(3x,A3)', advance='no') ' . '
            ! end if
            write(*, '(3x,f10.2)', advance='no') Jac(i,j)
        end do
        print *
    end do
  end subroutine print_sparsity_pattern

  subroutine output_jacobian(Jac)
    
    real(r8), intent(in) :: Jac(:, :)         ! Jacobian matrix
    integer :: i, j, N

    N = size(Jac, 1)
    open(unit=10, file='jacobian.txt', status='replace')
    
    write(10, '(A8)', advance='no') 'Row\Col'
    do j = 1, N
      write(10, '(i5)', advance='no') j
    end do
    write(10,*)

    do i = 1, N
      write(10, '(I5, A1)', advance='no') i, '|'
      do j = 1, N
        write(10, '(3x,f10.2)', advance='no') Jac(i,j)
      end do
      write(10,*)
    end do

    close(unit=10)
    print*, 'Jacobian matrix outputted to jacobian.txt'

  end subroutine output_jacobian

end module time_stepping_mod